# This library is free software: you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation, either
# version 3 of the License, or (at your option) any later version.

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library.  If not, see <http://www.gnu.org/licenses/> or <http://www.gnu.org/licenses/lgpl.txt>.

import socket
from . import ntlm
from ..py23 import get_moved_attr, import_moved
import inspect

urllib_request = import_moved('urllib2', 'urllib.request')
HTTPConnection = get_moved_attr('httplib', 'http.client', 'HTTPConnection')
HTTPSConnection = get_moved_attr('httplib', 'http.client', 'HTTPSConnection')
addinfourl = get_moved_attr('urllib', 'urllib.response', 'addinfourl')
URLError = get_moved_attr('urllib2', 'urllib.error', 'URLError')

def debug_output(*args, **kwargs):
    lineno = inspect.currentframe().f_back.f_lineno
    #print("debug at line ", lineno, ": ", *args, **kwargs)

class AbstractNtlmAuthHandler:

    def __init__(self, password_mgr=None):
        if password_mgr is None:
            password_mgr = urllib_request.HTTPPasswordMgr()
        self.passwd = password_mgr
        self.add_password = self.passwd.add_password

    def http_error_authentication_required(self, auth_header_field, req, fp, headers):
        auth_header_value_list = headers.get_all(auth_header_field)
        if auth_header_value_list:
            if any([hv.lower() == 'ntlm' for hv in auth_header_value_list]):
                fp.close()
                return self.retry_using_http_NTLM_auth(req, auth_header_field, None, headers)

    def retry_using_http_NTLM_auth(self, req, auth_header_field, realm, headers):
        user, pw = self.passwd.find_user_password(realm, req.get_full_url())
        debug_output(user, pw)
        if pw is not None:
            user_parts = user.split('\\', 1)
            if len(user_parts) == 1:
                UserName = user_parts[0]
                DomainName = ''
                type1_flags = ntlm.NTLM_ttype1_FLAGS & ~ntlm.NTLM_NegotiateOemDomainSupplied
            else:
                DomainName = user_parts[0].upper()
                UserName = user_parts[1]
                type1_flags = ntlm.NTLM_ttype1_FLAGS
            # ntlm secures a socket, so we must use the same socket for the complete handshake
            debug_output(hex(type1_flags))
            headers = dict(req.headers)
            headers.update(req.unredirected_hdrs)
            auth = 'NTLM %s' % ntlm.create_NTLM_NEGOTIATE_MESSAGE(user, type1_flags)
            if req.headers.get(self.auth_header, None) == auth:
                debug_output("no auth_header")
                return None
            headers[self.auth_header] = auth

            host = req.get_host()
            if not host:
                raise URLError('no host given')
            h = None
            if req.get_full_url().startswith('https://'):
                h = HTTPSConnection(host) # will parse host:port
            else:
                h = HTTPConnection(host) # will parse host:port
            # we must keep the connection because NTLM authenticates the connection, not single requests
            headers["Connection"] = "Keep-Alive"
            headers = dict((name.title(), val) for name, val in list(headers.items()))
            h.request(req.get_method(), req.get_selector(), req.data, headers)
            r = h.getresponse()
            r.begin()
            r._safe_read(int(r.getheader('content-length')))
            debug_output('data read')
            try:
                if r.getheader('set-cookie'):
                    # this is important for some web applications that store authentication-related info in cookies (it took a long time to figure out)
                    headers['Cookie'] = r.getheader('set-cookie')
                    debug_output('cookie: ', headers['Cookie'])
            except TypeError:
                debug_output('no cookie')
                pass
            r.fp = None # remove the reference to the socket, so that it can not be closed by the response object (we want to keep the socket open)
            auth_header_value = r.getheader(auth_header_field, None)
            debug_output(r.headers)
            debug_output(auth_header_field, ': ', auth_header_value)
            (ServerChallenge, NegotiateFlags) = ntlm.parse_NTLM_CHALLENGE_MESSAGE(auth_header_value[5:])
            debug_output('server c ', ServerChallenge, ' server flags ', hex(NegotiateFlags))
            auth = 'NTLM %s' % ntlm.create_NTLM_AUTHENTICATE_MESSAGE(ServerChallenge, UserName, DomainName, pw, NegotiateFlags)
            headers[self.auth_header] = auth
            debug_output('auth ', auth)
            headers["Connection"] = "Close"
            headers = dict((name.title(), val) for name, val in list(headers.items()))
            try:
                h.request(req.get_method(), req.get_selector(), req.data, headers)
                # none of the configured handlers are triggered, for example redirect-responses are not handled!
                response = h.getresponse()
                debug_output('data 3 read')
                def notimplemented():
                    raise NotImplementedError
                response.readline = notimplemented
                return addinfourl(response, response.msg, req.get_full_url())
            except socket.error as err:
                raise URLError(err)
        else:
            return None


class HTTPNtlmAuthHandler(AbstractNtlmAuthHandler, urllib_request.BaseHandler):

    auth_header = 'Authorization'
    handler_order = 480 # before Digest & Basic auth

    def http_error_401(self, req, fp, code, msg, headers):
        return self.http_error_authentication_required('www-authenticate', req, fp, headers)


class ProxyNtlmAuthHandler(AbstractNtlmAuthHandler, urllib_request.BaseHandler):
    """
        CAUTION: this class has NOT been tested at all!!!
        use at your own risk
    """
    auth_header = 'Proxy-authorization'
    handler_order = 480 # before Digest & Basic auth

    def http_error_407(self, req, fp, code, msg, headers):
        return self.http_error_authentication_required('proxy-authenticate', req, fp, headers)


if __name__ == "__main__":
    url = "http://ntlmprotectedserver/securedfile.html"
    user = "DOMAIN\\User"
    password = "Password"
    passman = urllib_request.HTTPPasswordMgrWithDefaultRealm()
    passman.add_password(None, url, user , password)
    auth_basic = urllib_request.HTTPBasicAuthHandler(passman)
    auth_digest = urllib_request.HTTPDigestAuthHandler(passman)
    auth_NTLM = HTTPNtlmAuthHandler(passman)

    # disable proxies (just for testing)
    proxy_handler = urllib_request.ProxyHandler({})

    opener = urllib_request.build_opener(proxy_handler, auth_NTLM) #, auth_digest, auth_basic)

    urllib_request.install_opener(opener)

    response = urllib_request.urlopen(url)
    print((response.read()))
